package baseclass.k_array;


import baseclass.a_sort.ArraysTool;

import java.util.Arrays;

/**
 * 使用partition进行topK。
 * 找topK先要找到第 k大元素,也就是index = k-1。
 * 如果换算成从小到大排序，也就是len - K +1 小。
 * <p>
 * 时间O(n)，空间：如果不允许该改原数组，要复制一份。O(n)。
 * <p>
 * 缺点：面对数据流，若存储需要巨大空间，不如堆。
 *
 * @date 2020/3/4 22:06
 */
public class Code06_TopK_Partition {

    public static int[] getTopK(int[] arr, int k) {
        if (arr == null || arr.length < k) return arr;
        int[] copy = Arrays.copyOf(arr, arr.length);
        //找到第 len - K +1 小，也就是index = len- k
//        int aim = select(copy, copy.length - k,0,arr.length-1);
        int aim = select(copy, copy.length - k);
        int[] res = new int[k];
        int p = 0;
        for (int i = 0; i < copy.length; i++) {
            if (copy[i] > aim) res[p++] = copy[i];
        }
        while (p < res.length) res[p++] = aim;
        return res;
    }

    /**
     * 方式1：递归版本，需要额外堆空间
     */
    private static int select(int[] copy,int targetIndex, int left,int right) {
        if(left >= right) return copy[left];
        int[] aimIndex = partition(copy, left, right,targetIndex);
        if(targetIndex >= aimIndex[0] && targetIndex <= aimIndex[1]) {
            return copy[aimIndex[0]];
        }else if(targetIndex < aimIndex[0]){
            return select(copy,targetIndex,left,aimIndex[0]-1);
        }else {
            return select(copy,targetIndex,aimIndex[1]+1,right);
        }
    }

    /**
     * 方式2：迭代版的select
     */
    private static int select(int[] copy,int targetIndex){
        int left = 0;
        int right = copy.length -1;
        int[] aimIndex = partition(copy,left,right,targetIndex);
        while (targetIndex < aimIndex[0] || targetIndex > aimIndex[1]){
            if(targetIndex < aimIndex[0]){
                aimIndex = partition(copy,left,aimIndex[0]-1,targetIndex);
            }else {
                aimIndex = partition(copy,aimIndex[1]+1,right,targetIndex);
            }
        }
        return copy[targetIndex];
    }
    private static int[] partition(int[] arr, int l, int r,int targetIndex) {
        int less = l - 1;
        int more = r + 1;
        int p = l;
        while (p < more) {
            if (arr[p] < arr[targetIndex]) {
                ArraysTool.swap(arr, ++less, p++);
            } else if (arr[p] > arr[targetIndex]) {
                ArraysTool.swap(arr, --more, p);
            } else {
                p++;
            }
        }
        return new int[]{less + 1, more - 1};
    }

    public static void main(String[] args) {
        int[] arr = {1,-1,5,2,4,342,342,51,1};
        int[] topK = getTopK(arr, 3);
        for (int i : topK) {
            System.out.println(i);
        }
    }
}
